package com.hunterai.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.hunterai.comfyui.ComfyUIService;
import com.hunterai.comfyui.ComfyUIWorkflowManager;
import com.hunterai.comfyui.ImageResponse;
import com.hunterai.comfyui.TaskStatusResponse;
import com.hunterai.dto.ReturnData;

import org.springframework.http.CacheControl;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.http.MediaType;

import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/api/cy")

public class ComfyUIProxyController {

    private final ComfyUIService comfyUIService;
    private final ComfyUIWorkflowManager workflowManager;
    private final ObjectMapper objectMapper;
    
    public ComfyUIProxyController(
            ComfyUIService comfyUIService, 
            ComfyUIWorkflowManager workflowManager,
            ObjectMapper objectMapper) {
        this.comfyUIService = comfyUIService;
        this.workflowManager = workflowManager;
        this.objectMapper = objectMapper;
    }
    
    /**
     * 获取所有可用工作流
     */
    @GetMapping("/workflows")
    public ResponseEntity<List<String>> getAvailableWorkflows() {
        List<String> workflows = workflowManager.getAvailableWorkflows();
        return ResponseEntity.ok(workflows);
    }
    
    /**
     * 获取指定工作流的参数配置
     */
    @GetMapping("/workflows/{workflowId}/params")
    public ResponseEntity<ReturnData<JsonNode>> getWorkflowParams(@PathVariable String workflowId) {
        JsonNode params = workflowManager.getWorkflowParams(workflowId);
        return ResponseEntity.ok(ReturnData.success(params));
    }

    /**
     * 获取指定工作流
     */
    @GetMapping("/workflows/{workflowId}")
    public ResponseEntity<JsonNode> getWorkflow(@PathVariable String workflowId) {
        JsonNode workflow = workflowManager.getWorkflow(workflowId);
        return ResponseEntity.ok(workflow);
    }
    
    /**
     * 生成图像API - 简化版
     */
    @PostMapping(
        value = "/generate/{workflowId}",
        consumes = MediaType.MULTIPART_FORM_DATA_VALUE
    )
    public ResponseEntity<?> generate(
        @PathVariable String workflowId,
        @RequestPart(value = "image", required = false) MultipartFile image,
        @RequestParam Map<String, Object> params
    ) {
        try {
            String imageName = null;
            if (image != null && !image.isEmpty()) {
                imageName = comfyUIService.imageUpload(image); // 上传并安全检查
            }
            // 传入 imageName（可为 null）
            JsonNode response = comfyUIService.generate(workflowId, imageName, params);
            return ResponseEntity.ok(response);
        } catch (IllegalArgumentException e) {
            return ResponseEntity.badRequest().body(Map.of("error", e.getMessage()));
        } catch (Exception e) {
            return ResponseEntity.badRequest().body(Map.of("error", "Generation failed: " + e.getMessage()));
        }
    }

    @GetMapping("/checktaskstatus")
    public ResponseEntity<?> checkTaskStatus(@RequestParam String promptId) {
        TaskStatusResponse response = comfyUIService.checkTaskStatus(promptId);
        
        if ("error".equals(response.getStatus())) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);
        } else if ("not_found".equals(response.getStatus())) {
            return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response);
        }
        
        return ResponseEntity.ok(response);
    }
    
    @GetMapping("/image")
    public ResponseEntity<?> getImage(
            @RequestParam String filename,
            @RequestParam(required = false, defaultValue = "") String subfolder,
            @RequestParam(required = false, defaultValue = "output") String type) {
        
        ImageResponse response = comfyUIService.getImage(filename, subfolder, type);
        
        if (response.getError() != null) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(Map.of("error", response.getError()));
        }
        
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(response.getMediaType());
        headers.setCacheControl(CacheControl.maxAge(1, TimeUnit.HOURS));
        
        return new ResponseEntity<>(response.getImageData(), headers, HttpStatus.OK);
    }
}
